import pandas as pd
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# 读取数据
df = pd.read_excel('./data/douban_movies_top250.xlsx')


# 提取国家和类型信息
def extract_countries_and_genres(info_str):
    parts = info_str.split('\xa0')
    countries = parts[0].strip().split(" ")
    genres = parts[-1].strip().split(" ")
    return countries, genres


df['countries'], df['genres'] = zip(*df['相关信息'].apply(extract_countries_and_genres))



# 获取所有不重复的国家和类型
all_countries = set([country for countries_list in df['countries'] for country in countries_list])
all_genres = set([genre for genres_list in df['genres'] for genre in genres_list])

print("--------------------")
print(all_countries)
print("--------------------")
print(all_genres)

# 独热编码
mlb_country = MultiLabelBinarizer()
mlb_genre = MultiLabelBinarizer()

country_matrix = mlb_country.fit_transform(df['countries'])
genre_matrix = mlb_genre.fit_transform(df['genres'])

# 合并国家和类型的独热编码矩阵
combined_matrix = np.hstack((country_matrix, genre_matrix))


def rec_by_country_and_genre(input_country, input_genre):
    # 为输入生成独热编码向量
    input_country = input_country.split(" ")
    input_genre = input_genre.split(" ")
    input_country_vec = mlb_country.transform(np.reshape(input_country, (1, -1)))
    input_genre_vec = mlb_genre.transform(np.reshape(input_genre, (1, -1)))
    input_vec = np.hstack((input_country_vec, input_genre_vec))

    # 计算相似度
    similarities = cosine_similarity(combined_matrix, input_vec)

    # 获取最相似的3部电影
    top3_indices = similarities.reshape((-1,)).argsort()[-3:][::-1]
    recommended_movies = df['影片中文名'].iloc[top3_indices]

    return recommended_movies.tolist()


def rec_by_movie(input_movie):
    # 检查输入电影是否在数据库中
    if input_movie not in df['影片中文名'].values:
        return "输入的电影不在数据库中，请输入有效的电影名。"

    # 获取输入电影的特征向量
    input_movie_index = df[df['影片中文名'] == input_movie].index[0]
    input_movie_vec = combined_matrix[input_movie_index]

    # 计算相似度
    similarities = cosine_similarity(combined_matrix, input_movie_vec.reshape(1, -1))

    # 获取最相似的3部电影（排除输入电影本身）
    top3_indices = similarities.reshape((-1,)).argsort()[-4:][::-1]  # 取前4个，因为第1个是电影本身
    top3_indices = top3_indices[top3_indices != input_movie_index]  # 排除电影本身
    recommended_movies = df['影片中文名'].iloc[top3_indices]

    return recommended_movies.tolist()
